import os
import pickle
import config
import data
import shutil
import glob


def unpack_data(problem_instance):
    """Map data.ProblemInstance to a list of Tensors, to process with map_fn."""
    if isinstance(problem_instance, data.ProblemInstance):
        return list(problem_instance)
    return problem_instance


def copy_checkpoint(checkpoint_path, global_step, accuracy):
    """Copies the checkpoint to a separate directory."""
    tmp_checkpoint_path = os.path.join(checkpoint_path, "tmp_best_checkpoint")
    best_checkpoint_path = os.path.join(checkpoint_path, "best_checkpoint")
    if _is_previous_accuracy_better(best_checkpoint_path, accuracy):
        print("Not copying the checkpoint: there is a better one from "
              "before a preemption.")
        return

    checkpoint_regex = os.path.join(checkpoint_path,
                                    "model.ckpt-{}.*".format(global_step))
    checkpoint_files = glob.glob(checkpoint_regex)
    graph_file = os.path.join(checkpoint_path, "graph.pbtxt")
    checkpoint_files.append(graph_file)

    _save_files_in_tmp_directory(tmp_checkpoint_path, checkpoint_files, accuracy)

    new_checkpoint_index_file = os.path.join(tmp_checkpoint_path, "checkpoint")
    f = os.open(new_checkpoint_index_file, mode="w")
    f.write("model_checkpoint_path: \"{}/model.ckpt-{}\"\n".format(best_checkpoint_path, global_step))

    # We first copy the better checkpoint to a temporary directory, and only
    # when it's created move it to avoid inconsistent state when job is preempted
    # when copying the checkpoint.
    if os.path.exists(best_checkpoint_path):
        shutil.rmtree(best_checkpoint_path)
    os.rename(tmp_checkpoint_path, best_checkpoint_path)
    print("Copied new best checkpoint with accuracy %.5f", accuracy)


def _save_files_in_tmp_directory(tmp_checkpoint_path, checkpoint_files,
                                 accuracy):
    """Saves the checkpoint files and accuracy in a temporary directory."""

    if os.path.exists(tmp_checkpoint_path):
        print("The temporary directory exists, because job was preempted "
              "before it managed to move it. We're removing it.")
        shutil.rmtree(tmp_checkpoint_path)
    os.mkdir(tmp_checkpoint_path)

    def dump_in_best_checkpoint_path(obj, filename):
        full_path = os.path.join(tmp_checkpoint_path, filename)
        f = open(full_path, mode="wb")
        pickle.dump(obj, f)

    for file_ in checkpoint_files:
        just_filename = file_.split("/")[-1]
        with open(file_) as f1, open(os.path.join(tmp_checkpoint_path, just_filename)) as f2:
            f2.write(f1.read())

    dump_in_best_checkpoint_path(config.get_inner_model_config(), "inner_config")
    dump_in_best_checkpoint_path(config.get_outer_model_config(), "outer_config")
    dump_in_best_checkpoint_path(accuracy, "accuracy")


def _is_previous_accuracy_better(best_checkpoint_path, accuracy):
    if not os.path.exists(best_checkpoint_path):
        return False

    previous_accuracy_file = os.path.join(best_checkpoint_path, "accuracy")
    f = open(previous_accuracy_file, mode="rb")
    previous_accuracy = pickle.load(f)

    return previous_accuracy > accuracy


def evaluate_and_average(session, tensor, num_estimates):
    tensor_value_estimates = [session.run(tensor) for _ in range(num_estimates)]
    average_tensor_value = sum(tensor_value_estimates) / num_estimates
    return average_tensor_value
